// Copyright (c) Aptos Foundation
// Licensed pursuant to the Innovation-Enabling Source Code License, available at https://github.com/aptos-labs/aptos-core/blob/main/LICENSE

use crate::{
    combinatorial_tests::{
        baseline::BaselineOutput,
        mock_executor::{MockEvent, MockTask},
        resource_tests::{
            create_executor_thread_pool, execute_block_parallel, get_gas_limit_variants,
        },
        types::{
            key_to_mock_module_id, KeyType, MockTransaction, TransactionGen, TransactionGenParams,
        },
    },
    task::ExecutorTask,
    txn_provider::default::DefaultTxnProvider,
};
use aptos_types::{state_store::MockStateView, transaction::AuxiliaryInfo};
use fail::FailScenario;
use move_core_types::language_storage::ModuleId;
use move_vm_types::module_id_interner::InternedModuleIdPool;
use proptest::{collection::vec, prelude::*, strategy::ValueTree, test_runner::TestRunner};
use test_case::test_matrix;

enum ModuleTestType {
    // All transactions publish modules, and all accesses are module reads.
    AllTransactionsAndAccesses,
    // All transactions publish modules, but some accesses are not module reads.
    AllTransactionsMixedAccesses,
    // Some transactions publish modules and contain module reads. Other
    // transactions do not publish modules and do not contain module reads.
    MixedTransactionsMixedAccesses,
}

fn execute_module_tests(
    universe_size: usize,
    transaction_count: usize,
    use_gas_limit: bool,
    blockstm_v2: bool,
    modules_test_type: ModuleTestType,
    num_executions: usize,
    num_random_generations: usize,
) where
    MockTask<KeyType<[u8; 32]>, MockEvent>:
        ExecutorTask<Txn = MockTransaction<KeyType<[u8; 32]>, MockEvent>>,
{
    let scenario = FailScenario::setup();
    assert!(fail::has_failpoints());
    fail::cfg("module_test", "return").unwrap();

    let executor_thread_pool = create_executor_thread_pool();
    let mut runner = TestRunner::default();

    let module_id_pool = InternedModuleIdPool::new();

    let gas_limits = get_gas_limit_variants(use_gas_limit, transaction_count);
    for gen_idx in 0..num_random_generations {
        // Generate universe
        let universe = vec(any::<[u8; 32]>(), universe_size)
            .new_tree(&mut runner)
            .expect("creating universe should succeed")
            .current();

        // Generate transactions based on parameters
        let transaction_strategy = match modules_test_type {
            ModuleTestType::AllTransactionsAndAccesses => vec(
                any_with::<TransactionGen<[u8; 32]>>(
                    TransactionGenParams::new_dynamic_modules_only(),
                ),
                transaction_count,
            ),
            ModuleTestType::AllTransactionsMixedAccesses
            | ModuleTestType::MixedTransactionsMixedAccesses => vec(
                any_with::<TransactionGen<[u8; 32]>>(
                    TransactionGenParams::new_dynamic_with_modules(),
                ),
                transaction_count,
            ),
        };

        let transaction_gen = transaction_strategy
            .new_tree(&mut runner)
            .expect("creating transactions should succeed")
            .current();

        // Convert transactions to use modules. For mixed transactions, we convert every
        // fifth transaction to use modules.
        let transactions: Vec<MockTransaction<KeyType<[u8; 32]>, MockEvent>> = transaction_gen
            .into_iter()
            .enumerate()
            .map(|(i, txn_gen)| {
                if i % 5 == 0
                    || !matches!(
                        modules_test_type,
                        ModuleTestType::MixedTransactionsMixedAccesses
                    )
                {
                    txn_gen.materialize_modules(&module_id_pool, &universe)
                } else {
                    txn_gen.materialize(&universe)
                }
            })
            .collect();

        let txn_provider = DefaultTxnProvider::new_without_info(transactions);
        let state_view = MockStateView::empty();

        // Generate all potential module IDs that could be used in the tests
        let all_module_ids = generate_all_potential_module_ids(&universe);

        // Run tests with fail point enabled to test the version metadata
        for exe_idx in 0..num_executions {
            for maybe_block_gas_limit in &gas_limits {
                if *maybe_block_gas_limit == Some(0) && (gen_idx > 0 || exe_idx > 0) {
                    // Run 0 gas limit test only once.
                    continue;
                }

                let output = execute_block_parallel::<
                    MockTransaction<KeyType<[u8; 32]>, MockEvent>,
                    MockStateView<KeyType<[u8; 32]>>,
                    DefaultTxnProvider<
                        MockTransaction<KeyType<[u8; 32]>, MockEvent>,
                        AuxiliaryInfo,
                    >,
                >(
                    executor_thread_pool.clone(),
                    *maybe_block_gas_limit,
                    &txn_provider,
                    &state_view,
                    Some(&all_module_ids),
                    blockstm_v2,
                );

                BaselineOutput::generate(txn_provider.get_txns(), *maybe_block_gas_limit)
                    .assert_parallel_output(&output);
            }
        }
    }
    scenario.teardown();
}

// Generate all potential module IDs that could be used in the tests
fn generate_all_potential_module_ids(universe: &[[u8; 32]]) -> Vec<ModuleId> {
    universe
        .iter()
        .map(|k| key_to_mock_module_id(&KeyType(*k), universe.len()))
        .collect()
}

// Test cases with various parameters
#[test_matrix(
    50, 100, [false, true], [false, true], [ModuleTestType::AllTransactionsAndAccesses, ModuleTestType::MixedTransactionsMixedAccesses,  ModuleTestType::AllTransactionsMixedAccesses], 10, 3; "module tests"
)]
#[test_matrix(
    10, 1000, [false, true], [false, true], [ModuleTestType::AllTransactionsAndAccesses, ModuleTestType::MixedTransactionsMixedAccesses,  ModuleTestType::AllTransactionsMixedAccesses], 5, 2; "contended module tests"
)]
fn module_transaction_tests(
    universe_size: usize,
    transaction_count: usize,
    use_gas_limit: bool,
    blockstm_v2: bool,
    modules_test_type: ModuleTestType,
    num_executions: usize,
    num_random_generations: usize,
) {
    execute_module_tests(
        universe_size,
        transaction_count,
        use_gas_limit,
        blockstm_v2,
        modules_test_type,
        num_executions,
        num_random_generations,
    );
}
